package com.archgeek.calcite;

import com.archgeek.calcite.Storage.DummyTable;
import org.apache.calcite.DataContext;
import org.apache.calcite.linq4j.AbstractEnumerable;
import org.apache.calcite.linq4j.Enumerable;
import org.apache.calcite.linq4j.Enumerator;
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.rel.type.RelDataTypeFactory;
import org.apache.calcite.rel.type.RelDataTypeFactory.FieldInfoBuilder;
import org.apache.calcite.schema.ScannableTable;
import org.apache.calcite.schema.Table;
import org.apache.calcite.schema.impl.AbstractTable;
import org.apache.calcite.sql.type.SqlTypeUtil;

import java.util.Arrays;

/**
 * @author pizhihui
 * @date 2024-05-30 22:54
 */
public class InMemoryTable extends AbstractTable implements ScannableTable {

    private String name;
    private Storage.DummyTable table;

    private RelDataType dataType;

    public InMemoryTable(String name) {
        this.name = name;
    }

    public InMemoryTable(String name, Storage.DummyTable table) {
        this.name = name;
        this.table = table;
    }

    @Override
    public Enumerable<Object[]> scan(DataContext root) {
        return new AbstractEnumerable<Object[]>() {
            @Override
            public Enumerator<Object[]> enumerator() {
                return new Enumerator<Object[]>() {
                    private int cur = 0;
                    @Override
                    public Object[] current() {
                        Object[] data = table.getData(cur++);
                        System.out.println("currrent:" + Arrays.toString(data));
                        return data;
                    }

                    @Override
                    public boolean moveNext() {
                        return cur < table.getRowCount();
                    }

                    @Override
                    public void reset() {

                    }

                    @Override
                    public void close() {

                    }
                };
            }
        };
    }

    @Override
    public RelDataType getRowType(RelDataTypeFactory typeFactory) {
        if (null == this.dataType) {
            FieldInfoBuilder builder = typeFactory.builder();
            this.table.getColumns().forEach(column -> {
                RelDataType type = typeFactory.createJavaType(column.getJavaClass());
                SqlTypeUtil.addCharsetAndCollation(type, typeFactory);
                builder.add(column.getName(), type);
            });
            this.dataType = typeFactory.createStructType(builder);
        }
        return this.dataType;
    }
}
